package karpiel.boguchwal.wojciech.przemek

import karpiel.boguchwal.wojciech.przemek.Ast.{Bruijn, PiType, PiValue}
import karpiel.boguchwal.wojciech.przemek.Cache.InnerCache
import karpiel.boguchwal.wojciech.przemek.RealCache.RealInnerCache
import karpiel.boguchwal.wojciech.przemek.Utils.usesBruins

import scala.collection.mutable


trait Cache {
  def getFor(ctx: Context): InnerCache
}

object Cache {

  trait InnerCache {
    def infer(ast: Ast, orElse: => Ast): Ast

    def normalize(ast: Ast, orElse: => Ast): Ast
  }

}


object Q {

  def usedBruijns(expr: Ast, ctx: Context) = {
    usesBruins(expr) ++ ctx.ctx.flatMap(e => usesBruins(e.typeT) ++ e.value.map(usesBruins(_)).getOrElse(Set()))
  }


  case class Squashed(expr: Ast, ctx: Context, oldToNew: Map[Int, Int])

  def squashContext(expr: Ast, ctx: Context, _usedBruijns: Set[Int]): Squashed = {
    val usedBruijns = _usedBruijns.toSeq.sorted
    val oldToNew = usedBruijns.zipWithIndex.toMap

    //teraz zaktualizować bruijny
    def withBruinUpdated(expr: Ast, minToAffect: Int): Ast = {
      expr match {
        case Bruijn(i, _) if i >= minToAffect =>
          Bruijn(oldToNew(i - minToAffect) + minToAffect)
        case Ast.PiValue(typeT, value) => PiValue(withBruinUpdated(typeT, minToAffect), withBruinUpdated(value, minToAffect + 1))
        case Ast.PiType(typeT, value) => PiType(withBruinUpdated(typeT, minToAffect), withBruinUpdated(value, minToAffect + 1))
        case q => Utils.applyToSubTrees(q, d => withBruinUpdated(d, minToAffect))
      }
    }

    val newCtx = Context(usedBruijns.map(ctx.ctx(_)).map { c =>
      ContextElement(withBruinUpdated(c.typeT, 0),
        c.value.map(v => withBruinUpdated(v, 0)),
        c.name)
    }.toList)


    val ast = withBruinUpdated(expr, 0)
    Squashed(ast, newCtx, oldToNew)
  }

  def restore(expr: Ast, oldToNew: Map[Int, Int]) = {
    val newToOld = oldToNew.map({ case (a, b) => b -> a }).toMap

    def modify(exprt: Ast, minToAffect: Int): Ast = {
      exprt match {
        case Bruijn(i, _) if i >= minToAffect => Bruijn(newToOld(i - minToAffect) + minToAffect)
        case Ast.PiValue(typeT, value) =>
          PiValue(modify(typeT, minToAffect), modify(value, minToAffect + 1))
        case Ast.PiType(typeT, value) =>
          PiType(modify(typeT, minToAffect), modify(value, minToAffect + 1))
        case q => Utils.applyToSubTrees(q, d => modify(d, minToAffect))

      }
    }

    val ast = modify(expr, 0)
    ast
  }
}

case class RealCache(cache: mutable.WeakHashMap[Context, InnerCache] = new mutable.WeakHashMap()) extends Cache {

  def handle(ctx: Context, expr: Ast) = {
    val ints = Q.usedBruijns(expr, ctx)
    val squashed = Q.squashContext(expr, ctx, ints)


    val restoredExpr = Q.restore(squashed.expr, squashed.oldToNew)
  }

  //TODO: cały czas trzeba haszować kontekst, dobrze jakby nie trzeba było
  def getFor(ctx: Context): InnerCache = {
    cache.getOrElseUpdate(ctx, RealInnerCache(ctx))
  }
}

object RealCache {

  case class RealInnerCache(ctx: Context, inferCache: mutable.WeakHashMap[Ast, Ast] = new mutable.WeakHashMap(), normalizationCache: mutable.WeakHashMap[Ast, Ast] = new mutable.WeakHashMap())
    extends Cache.InnerCache {


    def infer(ast: Ast, orElse: => Ast): Ast = {
      val ints = Q.usedBruijns(ast, ctx)
      val squashed = Q.squashContext(ast, ctx, ints)

      val restored = Q.restore(squashed.expr, squashed.oldToNew)


      inferCache.getOrElseUpdate(ast, orElse)
    }

    def normalize(ast: Ast, orElse: => Ast): Ast = normalizationCache.getOrElseUpdate(ast, orElse)
  }

}


object DummyCache extends Cache {
  override def getFor(ctx: Context): InnerCache = DummyInnerCache

  object DummyInnerCache extends InnerCache {
    override def infer(ast: Ast, orElse: => Ast): Ast = orElse

    override def normalize(ast: Ast, orElse: => Ast): Ast = orElse
  }

}